"""
@File  :make_shard.py
@Author:Xuelong Geng
@Date  :2024/6/7 17:02
@Desc  :
"""
import os

import tqdm
from gxl_ai_utils.utils import utils_file


def handle_line(input_line, res_wav_dict: dict, res_text_dict: dict):
    key, wav_path, txt_path = input_line.split()
    utils_file.logging_limit_print(key, wav_path, txt_path)
    # 完善path路径
    # txt提取， 英文转大写，标点去除
    wav_path = wav_path.replace("..", "/home/local_data2/lhma/Opensource")
    txt_path = txt_path.replace("..", "/home/local_data2/lhma/Opensource")
    utils_file.logging_limit_print(key, wav_path, txt_path)
    txt_line = utils_file.load_first_row_clean(txt_path)
    temp_list = txt_line.split()
    if len(temp_list) == 1:
        utils_file.logging_print("warning:txt_line=%s" % txt_line)
        return
    elif len(temp_list) >= 2:
        key, txt = temp_list[0], temp_list[1]
    utils_file.logging_limit_print(key, wav_path, txt)
    txt = utils_file.do_filter(txt)
    assert isinstance(txt, str)
    txt = txt.upper()
    utils_file.logging_limit_print(key, wav_path, txt)

    res_wav_dict[key] = wav_path
    res_text_dict[key] = txt


if __name__ == '__main__':
    input_tts_data_list_path = "/home/local_data2/lhma/Opensource/filelists/Basic_filelist.lst"
    lines_all = utils_file.load_list_file_clean(input_tts_data_list_path)
    utils_file.logging_print(f"lines_all={len(lines_all)}")


    def little_func(input_line_list, res_wav_dict, res_text_dict):
        for line_i in tqdm.tqdm(input_line_list, total=len(input_line_list)):
            handle_line(line_i, res_wav_dict, res_text_dict)


    runner = utils_file.GxlDynamicThreadPool()
    num_thread = 32
    list_list = utils_file.do_split_list(lines_all, num_thread)
    res_wav_dict = {}
    res_text_dict = {}
    for list_i in list_list:
        runner.add_task(little_func, [list_i, res_wav_dict, res_text_dict])
    runner.start()
    utils_file.logging_print("完成")
    output_dir = "/home/work_nfs14/xlgeng/asr_data_raw/wenetspeech4tts"
    utils_file.write_dict_to_scp(res_wav_dict, os.path.join(output_dir, "wav.scp"))
    utils_file.write_dict_to_scp(res_text_dict, os.path.join(output_dir, "text"))
    output_shard_dir = "/home/work_nfs14/xlgeng/asr_data_shard/wenetspeech4tts"
    # utils_file.do_make_shard_file(os.path.join(output_dir, "wav.scp"),os.path.join(output_dir, "text"), output_dir=output_shard_dir,num_threads=32)
